import torch
import cv2 as cv
from model.yolov5 import YoloV5
from model.decode import decode_box, scale_coords, plot_one_box
from utils.utils import letter_box_ractangle


model = YoloV5()
model.load_state_dict(torch.load('yolov5.pth'))
fourcc = cv.VideoWriter_fourcc(*'DIVX')
cap=cv.VideoCapture('D:/v_TaiChi_g03_c04.avi')
video = cv.VideoWriter('D:/video1.avi', fourcc, 25, (320, 240), True)
model.eval()
print('Network loading complete.')
# t1=0
while cap.isOpened():
    # print(time.time()-t1)
    # t1=time.time()
    ret, image_origin = cap.read()
    if not ret:
        break
    image = letter_box_ractangle(image_origin,320)
    image = image[:, :, ::-1].transpose(2, 0, 1)/255
    with torch.no_grad():
        image = torch.tensor(image, dtype=torch.float32).unsqueeze(0)
        predict = model(image)
        predict = decode_box(predict)
    det = predict[0]
    det[:, :4] = scale_coords(image.shape[2:], det[:, :4], image_origin.shape).round()
    for *xyxy, conf, cls in reversed(det):
        label = f'{int(cls)} {conf:.4f}'
        plot_one_box(xyxy, image_origin, label=label, color=(0, 0, 255), line_thickness=1)
    cv.imshow('test', image_origin)
    cv.waitKey(1)
    video.write(image_origin)
cap.release()
video.release()
